from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

# analysis of phase space

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal

# 全局变量，用于存储和图像显示线程交互的数据
class imgview:
    global_image = None
    imgview_exit = False
    trajectory = []
    focus_i = 0
    traj_i = 0

imgview_data = imgview()

# 定义一个函数，用于在独立线程中显示图像
def show_image():
    grid_size_display = 20
    while not imgview_data.imgview_exit:
        # 检查全局变量是否有图像
        if imgview_data.global_image is not None:
            img = np.copy(imgview_data.global_image)
            state_x = imgview_data.trajectory[imgview_data.traj_i][0]
            state_y = imgview_data.trajectory[imgview_data.traj_i][1]
            cv2.circle(img, (state_y * grid_size_display + int(grid_size_display/2), state_x * grid_size_display + int(grid_size_display/2)), 7, (0, 0, 255), -1, cv2.LINE_AA)
            # 显示图像
            cv2.imshow("Image", img)
            key = cv2.waitKey(1)
            if key == ord('a'):
                imgview_data.focus_i -= 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
            elif key == ord('d'):
                imgview_data.focus_i += 1
                print("imgview_data.focus_i: ", imgview_data.focus_i)
        else:
            # 图像还未产生，等待100毫秒
            time.sleep(0.1)

def rnn_run(x, mat_intr, bias1):
    intr_vector = np.dot(x, mat_intr) + bias1
    intr_vector = np.tanh(intr_vector)
    return intr_vector

# 计算雅克比矩阵
@jax.jit
def rnn_run_vector(x, mat_intr, bias1):
    intr_vector = jnp.dot(x, mat_intr) + bias1
    intr_vector = jnp.tanh(intr_vector)
    return intr_vector - x
jacobian_fun = jax.jacrev(rnn_run_vector)
jacobian_vmap = jax.jit(jax.vmap(jacobian_fun, in_axes=(0,None,None)))

@jax.jit
def vector_field_taylor_expansion(x, x0, jacobian, mat_intr, bias1):
    vector = jnp.dot(jacobian, x - x0) + rnn_run_vector(x0, mat_intr, bias1)
    return vector
vector_field_taylor_expansion_vmap = jax.jit(jax.vmap(vector_field_taylor_expansion, in_axes=(0,0,None,None,None)))

# run linear dynamical system using vector_field_taylor_expansion
def run_LDS_one_step(x, x0, dt, jacobian, mat_intr, bias1):

    x_dot = vector_field_taylor_expansion(x, x0, jacobian, mat_intr, bias1)
    x1 = x + x_dot * dt

    return x1.copy()

def compute_linearized_err_1X(sample_point, mat_intr, bias1):

    jacobian_of_x = jacobian_fun(jnp.array(sample_point), mat_intr, bias1)

    sample_vector = rnn_run_vector(sample_point, mat_intr, bias1)
    sample_vector_len = np.linalg.norm(sample_vector)

    initial_state = sample_point
    current_state = initial_state.copy()
    current_state0 = current_state.copy()
    traj_len = 0

    next_state = rnn_run(initial_state, mat_intr, bias1)

    while traj_len < sample_vector_len:
        current_state0 = current_state.copy()
        current_state = run_LDS_one_step(current_state, initial_state, 0.01, jacobian_of_x, mat_intr, bias1)
        delta_state = current_state - initial_state
        delta_state_len = np.linalg.norm(delta_state)
        traj_len += delta_state_len

    diff_next_state_current_state = np.linalg.norm(next_state - current_state0)

    return diff_next_state_current_state, np.corrcoef(next_state, current_state0)[0, 1]

def compute_eig(sample_point, mat_intr, bias1):
    jacobian_of_x = jacobian_fun(jnp.array(sample_point), mat_intr, bias1)
    eigenvalues = np.linalg.eigvals(jacobian_of_x)
    if any(np.real(eigenvalues) > 0):
        return False, eigenvalues
    else:
        return True, eigenvalues

def compute_limit_point(sample_point, mat_intr, bias1):

    jacobian_of_x = jacobian_fun(jnp.array(sample_point), mat_intr, bias1)

    initial_state = sample_point
    current_state = initial_state.copy()

    for i in range(1000):
        current_state = run_LDS_one_step(current_state, initial_state, 0.01, jacobian_of_x, mat_intr, bias1)

    return current_state

def main():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)
    
    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 5000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)
    rnn_state_old = rnn_state.copy()
    diff = jnp.abs(rnn_state - rnn_state_old)
    rnn_state_old = rnn_state.copy()
    diff_norm = jnp.linalg.norm(diff, axis=1)
    diff_norm_old = diff_norm.copy()
    norm_std = diff_norm.copy()

    rnn_state_init = rnn_state.copy()

    obs_zero = jnp.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] for i in range(n_samples)])

    rnn_state_trajectory = []

    for t in range(rpl_config.life_duration):

        if t == rpl_config.probe_point:
            rnn_state_init = rnn_state.copy()

        progress_bar(t, rpl_config.life_duration)

        """ model forward 
        """
        rnn_state, y1 = model_forward(params, rnn_state, obs_zero, model)
        diff = jnp.abs(rnn_state - rnn_state_old)
        rnn_state_old = rnn_state.copy()
        diff_norm = jnp.linalg.norm(diff, axis=1)
        norm_std = 0.4 * jnp.abs(diff_norm - diff_norm_old) + 0.6 * norm_std
        diff_norm_old = diff_norm.copy()

        rnn_state_trajectory.append(np.array(rnn_state).copy())
            
    print(rnn_state.shape)
    print(norm_std.shape)
    rnn_state_np = np.array(rnn_state)

    # 将 rnn_state_trajectory 展开成 rnn_state_np 的形状
    rnn_state_trajectory_np = np.array(rnn_state_trajectory)
    # 交换 rnn_state_trajectory_np 的第一维和第二维
    rnn_state_trajectory_np = np.swapaxes(rnn_state_trajectory_np, 0, 1)
    print("shape of rnn_state_trajectory_np: ", rnn_state_trajectory_np.shape)

    # 计算 rnn_state_np 中，每个样本和其 前一个样本之间的差向量的模长
    @jit
    def compute_diff_norm(rnn_state_np):
        diff = rnn_state_np[1:] - rnn_state_np[:-1]
        diff_norm = jnp.linalg.norm(diff, axis=1)
        return diff_norm
    compute_diff_norm_vmap = jax.vmap(compute_diff_norm)

    @jit
    def compute_diff_norm1(rnn_state_np):
        diff_norm = rnn_state_np[1:] - rnn_state_np[:-1]
        return diff_norm
    compute_diff_norm_vmap1 = jax.vmap(compute_diff_norm1)

    rnn_state_trajectory_velocity = compute_diff_norm_vmap(rnn_state_trajectory_np)
    print("shape of rnn_state_trajectory_velocity: ", rnn_state_trajectory_velocity.shape)

    rnn_state_trajectory_velocity_delta = compute_diff_norm_vmap1(rnn_state_trajectory_velocity)
    print("shape of rnn_state_trajectory_velocity_delta: ", rnn_state_trajectory_velocity_delta.shape)
    rnn_state_trajectory_velocity = rnn_state_trajectory_velocity[:,1:]

    rnn_state_trajectory_velocity = np.array(rnn_state_trajectory_velocity)
    rnn_state_trajectory_velocity_delta = np.array(rnn_state_trajectory_velocity_delta)

    rnn_state_trajectory_velocity_damping_ratio = rnn_state_trajectory_velocity.astype(float) / (rnn_state_trajectory_velocity_delta.astype(float)+0.0001)
    print("shape of rnn_state_trajectory_velocity_damping_ratio: ", rnn_state_trajectory_velocity_damping_ratio.shape)
    # 将 rnn_state_trajectory_velocity_damping_ratio 中无穷大的值替换为 0
    rnn_state_trajectory_velocity_damping_ratio[np.isinf(rnn_state_trajectory_velocity_damping_ratio)] = 0
    # 将 rnn_state_trajectory_velocity_damping_ratio 中的 nan 值替换为 0
    rnn_state_trajectory_velocity_damping_ratio[np.isnan(rnn_state_trajectory_velocity_damping_ratio)] = 0

    # 将 rnn_state_trajectory_velocity_damping_ratio rnn_state_trajectory_velocity 小于0.1的值替换为 0
    rnn_state_trajectory_velocity_damping_ratio[rnn_state_trajectory_velocity < 0.1] = 0
    rnn_state_trajectory_velocity_damping_ratio[rnn_state_trajectory_velocity_damping_ratio>100]=0
    rnn_state_trajectory_velocity_damping_ratio[rnn_state_trajectory_velocity_damping_ratio<-100]=0

    rnn_state_trajectory_velocity_damping_ratio = np.swapaxes(rnn_state_trajectory_velocity_damping_ratio, 0, 1)

    print("shape of rnn_state_trajectory_velocity_damping_ratio: ", rnn_state_trajectory_velocity_damping_ratio.shape)

    positive_dampping_ratios = []
    negative_dampping_ratios = []
    for i in range(rnn_state_trajectory_velocity_damping_ratio.shape[0]):
        ratios = rnn_state_trajectory_velocity_damping_ratio[i]
        # 统计 ratios 中大于0的值的个数以及小于0的值的个数
        positive_dampping_ratios.append(np.sum(ratios>0))
        negative_dampping_ratios.append(np.sum(ratios<0))
    
    positive_dampping_ratios = np.array(positive_dampping_ratios)
    negative_dampping_ratios = np.array(negative_dampping_ratios)

    # 将 positive_dampping_ratios 和 negative_dampping_ratios 绘制到同一个图表中，用plot
    plt.plot(positive_dampping_ratios, color='red')
    plt.plot(negative_dampping_ratios, color='blue')
    # 横轴取名为 "time step"，纵轴取名为 "number of positive/negative damping ratios"
    plt.xlabel("time step")
    plt.ylabel("number of positive/negative damping ratios")
    # 显示图例
    plt.legend(["positive", "negative"])

    plt.show()

    # # 将 rnn_state_trajectory_velocity 显示为 3D 曲面
    # # Generate some sample data
    # x = np.arange(rnn_state_trajectory_velocity_damping_ratio.shape[1])
    # y = np.arange(rnn_state_trajectory_velocity_damping_ratio.shape[0])
    # X, Y = np.meshgrid(x, y)
    # Z = rnn_state_trajectory_velocity_damping_ratio

    # # Create a new figure
    # fig = plt.figure()

    # # Create a 3D axes object
    # ax = fig.add_subplot(111, projection='3d')

    # # Plot the surface
    # ax.plot_surface(X, Y, Z)

    # # Set the labels for the axes
    # ax.set_xlabel('X')
    # ax.set_ylabel('Y')
    # ax.set_zlabel('Z')

    # # Show the plot
    # plt.show()

if __name__ == "__main__":
    main()